package com.zzy.chapter4.iterator;

import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.client.RestTemplate;

import java.util.*;
import java.util.stream.Collectors;

public class EsQueryIterator implements Iterator<Map<String, Object>>{

    // 记录当前 cursor 分页
    private String cursor;

    // 记录查询到的columns
    private List<String> columns;

    // 将ES SQL RestAPI 返回值记录
    private Iterator<Map<String, Object>> iterator;

    private RestTemplate restTemplate = new RestTemplate();

    private final String ES_URL = "http://localhost:19200/_sql?format=json";

    public EsQueryIterator(String query,Long fetchSize) {
        EsResponseData esResponseData = restTemplate.postForObject(ES_URL, new EsSqlQuery(query, fetchSize), EsResponseData.class);
        assert esResponseData != null;
        this.columns = esResponseData.getColumns()
                .stream()
                .map(x -> x.get("name"))
                .collect(Collectors.toList());
        this.iterator = convert(columns, esResponseData).iterator();
    }

    @Override
    public boolean hasNext() {
        return iterator.hasNext() || scrollNext();
    }

    private boolean scrollNext() {
        if(iterator == null || this.cursor == null) {
            return false;
        }
        EsResponseData esResponseData = restTemplate.postForObject(ES_URL, new EsSqlQuery(cursor), EsResponseData.class);
        this.cursor = esResponseData.getCursor();
        this.iterator = convert(columns, esResponseData).iterator();
        return iterator.hasNext();
    }

    @Override
    public Map<String, Object> next() {
        return iterator.next();
    }

    private List<Map<String, Object>> convert(List<String> columns, EsResponseData esResponseData) {
        List<Map<String, Object>> results = new ArrayList<>();
        for (List<Object> row : esResponseData.getRows()) {
            Map<String, Object> map = new HashMap<>();
            for (int i = 0; i < columns.size(); i++) {
                map.put(columns.get(i), row.get(i));
            }
            results.add(map);
        }
        return results;
    }

}
